"""
多文档问答机器人 - 基于 RAG (Retrieval-Augmented Generation) 技术

功能特点：
- 支持多种文档格式：PDF、Word、TXT
- 基于 Streamlit 的友好聊天界面
- 使用 Ollama 本地模型进行嵌入和生成
- 具备对话记忆功能

安装依赖：
pip install PyPDF2 docx streamlit langchain ollama python-dotenv

运行命令：
streamlit run .\langchain\llm_langchain_14_RAG多文档问答机器人.py
"""

import docx
import streamlit as st
from PyPDF2 import PdfReader
from dotenv import load_dotenv
# ConversationalRetrievalChain 用于构建 RAG 聊天链：对话记忆 + 文本检索 + LLM回答
from langchain.chains.conversational_retrieval.base import ConversationalRetrievalChain
from langchain_community.vectorstores import FAISS
from langchain_core.documents import Document
from langchain_ollama import OllamaEmbeddings, ChatOllama
from langchain_text_splitters import RecursiveCharacterTextSplitter

# 加载环境变量（用于 OpenAI 等云端模型，当前使用本地 Ollama）
load_dotenv()

# 设置 Streamlit 应用页面属性
st.set_page_config(page_title="智能文档问答助手", page_icon="🤖", layout="wide")

# 注入 CSS 用于美化聊天界面
st.markdown(
    """<style>
    /* 聊天消息基础样式 */
    .chat-message {
        padding: 1.5rem; 
        border-radius: 0.5rem; 
        margin-bottom: 1rem; 
        display: flex
    }
    /* 用户消息样式 - 深色背景 */
    .chat-message.user {
        background-color: #2b313e
    }
    /* 机器人消息样式 - 稍浅背景 */
    .chat-message.bot {
        background-color: #475063
    }
    /* 头像容器 */
    .chat-message .avatar {
        width: 20%;
    }
    /* 头像图片样式 */
    .chat-message .avatar img {
        max-width: 78px;
        max-height: 78px;
        border-radius: 50%;
        object-fit: cover;
    }
    /* 消息内容区域 */
    .chat-message .message {
        width: 80%;
        padding: 0 1.5rem;
        color: #fff;
    }
    /* 隐藏部署按钮和菜单 */
    .stDeployButton {
        visibility: hidden;
    }
    #MainMenu {visibility: hidden;}
    footer {visibility: hidden;}
    
    /* 调整主容器内边距 */
    .block-container {
        padding: 2rem 4rem 2rem 4rem;
    }
    
    /* 侧边栏样式调整 */
    .st-emotion-cache-16txtl3 {
        padding: 3rem 1.5rem;
    }
    </style>
    """,
    unsafe_allow_html=True,
)

# 定义机器人回复消息 HTML 模板
bot_template = """
<div class="chat-message bot">
    <div class="avatar">
        <img src="https://cdn.icon-icons.com/icons2/1371/PNG/512/robot02_90810.png" style="max-height: 78px; max-width: 78px; border-radius: 50%; object-fit: cover;">
    </div>
    <div class="message">{{MSG}}</div>
</div>
"""

# 定义用户消息 HTML 模板
user_template = """
<div class="chat-message user">
    <div class="avatar">
        <img src="https://www.shareicon.net/data/512x512/2015/09/18/103160_man_512x512.png" >
    </div>    
    <div class="message">{{MSG}}</div>
</div>
"""


def get_file_text(uploaded_files):
    """
    读取上传的多种格式文档并转换为 Document 对象列表

    Args:
        uploaded_files: Streamlit 上传的文件对象列表

    Returns:
        list: Document 对象列表，包含文本内容和元数据

    支持格式：
    - PDF: 逐页提取文本
    - Word: 逐段提取文本
    - TXT: 直接读取文本内容
    """
    docs = []
    for document in uploaded_files:
        if document.type == "application/pdf":
            # 处理 PDF 文件
            pdf_reader = PdfReader(document)
            for idx, page in enumerate(pdf_reader.pages):
                # 提取 PDF 每一页文本，并记录来源信息
                docs.append(
                    Document(
                        page_content=page.extract_text(),
                        metadata={"source": f"{document.name} 第{idx + 1}页"}
                    )
                )
        elif document.type == "application/vnd.openxmlformats-officedocument.wordprocessingml.document":
            # 处理 Word (.docx) 文档
            doc = docx.Document(document)
            for idx, paragraph in enumerate(doc.paragraphs):
                if paragraph.text.strip():  # 跳过空段落
                    docs.append(
                        Document(
                            page_content=paragraph.text,
                            metadata={"source": f"{document.name} 第{idx + 1}段"}
                        )
                    )
        elif document.type == "text/plain":
            # 处理纯文本文件
            text = document.getvalue().decode("utf-8")
            docs.append(
                Document(
                    page_content=text,
                    metadata={"source": f"{document.name}"}
                )
            )

    st.success(f"成功读取 {len(docs)} 个文档片段")
    return docs


def get_text_chunks(docs):
    """
    使用递归文本切分器将文档切分为适合处理的文本块

    Args:
        docs: Document 对象列表

    Returns:
        list: 切分后的 Document 对象列表

    为什么需要切分：
    - 适应模型输入长度限制
    - 提高检索精度
    - 避免上下文过长
    """
    # 创建递归字符文本切分器
    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=500,  # 每个文本块的最大字符数
        chunk_overlap=10  # 块之间的重叠字符数，防止语义断层
    )
    docs_chunks = text_splitter.split_documents(docs)
    st.info(f"文档切分完成：{len(docs)} → {len(docs_chunks)} 个文本块")
    return docs_chunks


def get_vectorstore(docs_chunks):
    """
    将文本块转换为向量并存储到 FAISS 向量数据库中

    Args:
        docs_chunks: 切分后的 Document 对象列表

    Returns:
        FAISS: 配置好的向量数据库实例

    流程：
    1. 初始化嵌入模型（将文本转换为向量）
    2. 将文档块转换为向量
    3. 构建 FAISS 向量索引
    """
    # 初始化 Ollama 嵌入模型（本地部署）
    embeddings = OllamaEmbeddings(base_url="http://127.0.0.1:11434", model="nomic-embed-text:latest")

    # 备用选项：使用 HuggingFace 嵌入模型
    # embeddings = HuggingFaceEmbeddings(
    #     model_name="D:/models/BAAI_bge-small-zh-v1.5",
    #     model_kwargs={'device': 'cpu'}
    # )

    # 创建向量数据库
    with st.spinner("正在构建向量数据库..."):
        vectorstore = FAISS.from_documents(
            documents=docs_chunks,
            embedding=embeddings
        )

    st.success("向量数据库构建完成！")
    return vectorstore


def get_conversation_chain(vectorstore):
    """
    构建对话式 RAG 检索链

    Args:
        vectorstore: 向量数据库实例

    Returns:
        ConversationalRetrievalChain: 配置好的对话链

    功能：
    - 输入用户问题
    - 检索相关文档片段
    - 结合对话历史生成回答
    - 返回答案和参考来源
    """
    # 初始化聊天模型（本地 Ollama）
    llm = ChatOllama(temperature=0, model="deepseek-r1:1.5b", base_url="http://127.0.0.1:11434")

    # 备用选项：使用 OpenAI 模型
    # llm = ChatOpenAI(
    #     model="Qwen/Qwen2.5-VL-72B-Instruct",
    #     streaming=True,
    #     temperature=0
    # )

    # 构建对话检索链
    conversation_chain = ConversationalRetrievalChain.from_llm(
        llm=llm,  # 语言模型
        retriever=vectorstore.as_retriever(),  # 检索器
        return_source_documents=True,  # 返回参考文档
        verbose=True  # 显示详细日志
    )

    st.success("对话链构建完成！")
    return conversation_chain


def handle_userinput(user_question):
    """
    处理用户输入问题，与 RAG 链交互，并显示回答及参考文档

    Args:
        user_question: 用户输入的问题文本
    """
    # 获取当前对话历史
    chat_history = st.session_state.chat_history

    # 调用对话链获取回答
    with st.spinner("正在思考..."):
        response = st.session_state.conversation(
            {
                "question": user_question,
                "chat_history": chat_history
            }
        )

    # 更新对话历史（记录用户问题和AI回答）
    st.session_state.chat_history.append(("user", user_question))
    st.session_state.chat_history.append(("assistant", response["answer"]))

    # 展示用户输入
    st.write(
        user_template.replace("{{MSG}}", user_question),
        unsafe_allow_html=True
    )

    # 处理并展示参考来源
    sources = response["source_documents"]
    if sources:
        # 提取所有来源的元数据信息
        source_names = [doc.metadata["source"] for doc in sources]
        sources_text = "\n\n".join(set(source_names))  # 去重
        sources_display = f"\n\n📚 参考来源：\n{sources_text}"
    else:
        sources_display = "\n\n⚠️ 未找到相关参考文档"

    # 显示AI回复和来源
    ai_response = st.session_state.chat_history[-1][1]  # 最新AI回复
    st.write(
        bot_template.replace("{{MSG}}", ai_response + sources_display),
        unsafe_allow_html=True
    )


def show_history():
    """
    展示完整的对话历史记录

    显示规则：
    - 偶数索引：用户消息
    - 奇数索引：AI回复
    """
    chat_history = st.session_state.chat_history
    for i, message in enumerate(chat_history):
        if i % 2 == 0:  # 用户消息
            st.write(
                user_template.replace("{{MSG}}", message[1]),
                unsafe_allow_html=True
            )
        else:  # AI回复
            st.write(
                bot_template.replace("{{MSG}}", message[1]),
                unsafe_allow_html=True
            )


def clear_chat_history():
    """清空对话历史记录"""
    st.session_state.chat_history = []
    st.success("对话历史已清空！")


def main():
    """
    主函数：程序入口点

    主要功能：
    1. 初始化界面和会话状态
    2. 处理文档上传和向量化
    3. 管理用户交互和对话显示
    4. 维护对话历史
    """
    # 页面顶部标题
    st.header("📚 智能文档问答助手")
    st.caption("上传文档后即可开始对话问答")

    # 初始化会话状态
    if "conversation" not in st.session_state:
        st.session_state.conversation = None
    if "chat_history" not in st.session_state:
        st.session_state.chat_history = []

    # 左侧边栏 - 文档管理区域
    with st.sidebar:
        st.title("📋 文档管理")

        # 文件上传器
        uploaded_files = st.file_uploader(
            "选择文件",
            type=["txt", "pdf", "doc", "docx"],
            accept_multiple_files=True,
            help="支持 TXT、PDF、Word 格式文档"
        )

        # 文档处理按钮
        if st.button(
                "🚀 处理文档",
                use_container_width=True,
                type="primary"
        ):
            if uploaded_files:
                with st.status("文档处理中...", expanded=True) as status:
                    st.write("📖 读取文档内容...")
                    docs = get_file_text(uploaded_files)

                    st.write("✂️ 切分文档...")
                    docs_chunks = get_text_chunks(docs)

                    st.write("🔢 构建向量数据库...")
                    vectorstore = get_vectorstore(docs_chunks)

                    st.write("🔗 创建对话链...")
                    st.session_state.conversation = get_conversation_chain(vectorstore)

                    status.update(label="文档处理完成！", state="complete")
            else:
                st.warning("请先上传文件！")

        # 清空对话历史按钮
        if st.session_state.chat_history:
            st.button(
                "🗑️ 清空对话",
                on_click=clear_chat_history,
                use_container_width=True
            )

    # 主界面 - 聊天区域
    with st.container():
        # 用户输入框
        user_question = st.chat_input("💭 请输入您的问题...")

        # 显示当前对话状态
        if st.session_state.conversation:
            st.success("✅ 系统就绪，可以开始提问！")
        else:
            st.info("📁 请先在左侧上传并处理文档")

    # 聊天内容展示区域
    with st.container():
        show_history()

        # 处理用户提问
        if user_question:
            if st.session_state.conversation is not None:
                handle_userinput(user_question)
            else:
                st.warning("请先上传并处理文档！")


if __name__ == "__main__":
    main()
